/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "add_vector_tiling.h"
#include "data_utils.h"

#include "acl/acl.h"
#include "aclrtlaunch_add_vector.h"
/*
获取一些硬件平台的信息，来支撑Tiling的计算，比如获取硬件平台的核数等信息。
PlatformAscendCManager类提供获取平台信息的功能：通过该类的GetInstance方法可以获取一个PlatformAscendC类的指针，再通过该指针获取硬件平台相关信息
*/
#include "tiling/platform/platform_ascendc.h"
#include "shape_config.h"

const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    
    size_t inputByteSize = SHAPE[0] * SHAPE[1] * sizeof(uint16_t);
    size_t outputByteSize = SHAPE[0] * SHAPE[1] * sizeof(uint16_t);

    auto platform_ascendc = platform_ascendc::PlatformAscendCManager::GetInstance();
    //获取当前硬件平台的核数。若AI Core的架构为Cube、Vector分离架构，返回AI Core上的Vector核数；非分离架构返回AI Core的核数。
    uint32_t coreNum = platform_ascendc->GetCoreNum(); 
    printf("可用核心数量 coreNum : %u\n", coreNum);
    
    // 获取UB内存大小
    uint64_t ubSize = 0;
    platform_ascendc->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    printf("UB memory size ubSize : %lu Bytes\n", ubSize);
    
    
    // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
    uint32_t inputNum = SHAPE[0] * SHAPE[1];
    printf("输入的元素/标量数量 inputNum : %u\n", inputNum);
    uint32_t typeLength = sizeof(uint16_t);
    
    uint32_t inputLength = inputNum * typeLength;
    printf("输入的元素/标量大小 typeLength : %u Bytes\n", typeLength);
    uint32_t inputBytes = inputLength / inputNum;
    printf("输入的大小(未对齐前) inputLength : %u Bytes\n\n", inputLength);

    // There are a total of 3 shared UB spaces in the input and output. If it's int8, there are 2 more TBUFs
    uint32_t ubDataNumber = (inputBytes == 1) ? 5 : 3;
    // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;
    printf("本算子UB可容纳的一次处理计算的块数（输入） tileBlockNum : %u\n", tileBlockNum);
    printf("本算子UB可容纳的一次处理计算的标量数量（输入） tileDataNum : %u\n\n", tileDataNum);

    // Input data for 32B alignment
    uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
    // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
    coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
    coreNum = (coreNum >= 1) ? coreNum : 1;
    coreNum = 1;
    uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;
    uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;
    printf("输入对齐到32B后的大小 inputLengthAlgin32 : %u Bytes\n", inputLengthAlgin32);
    printf("分配计算的核心数量 coreNum : %u\n", coreNum);
    printf("每个核心分配的输入块数 everyCoreInputBlockNum : %u\n", everyCoreInputBlockNum);
    printf("尾块数量（大核数量） tailBlockNum : %u\n\n", tailBlockNum);
    
    // Small chunks are calculated and sliced several times using the number of data on each core
    uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
    // Tail block calculation for small chunks of data
    uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
    smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    
    // Print statements for small chunks calculation
    printf("每个小核分配的输入元素/标量数量 smallCoreDataNum : %u \n", smallCoreDataNum);
    printf("每个小核从UB搬运输入的次数（往下取整） smallTileNum : %u\n", smallTileNum);
    printf("每个小核从UB搬运输入的次数 finalSmallTileNum : %u\n", finalSmallTileNum);
    printf("每个小核从UB搬运输入的剩余元素/标量数量 smallTailDataNum : %u \n\n", smallTailDataNum);
    
    // The total length of a large block of data is 32B larger than that of a small block of data
    everyCoreInputBlockNum += 1;
    uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
    uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
    bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;
    
    // Print statements for large chunks calculation
    printf("每个大核分配的输入元素/标量数量 bigCoreDataNum : %u \n", bigCoreDataNum);
    printf("每个大核从UB搬运输入的次数（往下取整） bigTileNum : %u\n", bigTileNum);
    printf("每个大核从UB搬运输入的次数 finalBigTileNum : %u\n", finalBigTileNum);
    printf("每个大核从UB搬运输入的剩余元素/标量数量 bigTailDataNum : %u \n", bigTailDataNum);

    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    AddVectorTilingData *tiling;
    uint8_t *xHost, *yHost, *zHost;
    uint8_t *xDevice, *yDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), sizeof(AddVectorTilingData)));
    tiling->smallCoreDataNum = smallCoreDataNum;
    tiling->bigCoreDataNum = bigCoreDataNum;
    tiling->finalBigTileNum = finalBigTileNum;
    tiling->finalSmallTileNum = finalSmallTileNum;
    tiling->tileDataNum = tileDataNum;
    tiling->smallTailDataNum = smallTailDataNum;
    tiling->bigTailDataNum = bigTailDataNum;
    tiling->tailBlockNum = tailBlockNum;

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);
    ReadFile("./input/input_y.bin", inputByteSize, yHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(yDevice, inputByteSize, yHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    blockDim = coreNum;
    
    ACLRT_LAUNCH_KERNEL(add_vector)(blockDim, stream, xDevice, yDevice, zDevice, tiling);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, outputByteSize);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());

    return 0;
}
